from pathlib import Path
from typing import NamedTuple
from queue import Queue


W, H = 1920, 1080
TYPE_DICT = {
    "chicken": 0
}


class GSSDData(NamedTuple):
    frame_id: str
    unused: str
    x: str
    y: str
    w: str
    h: str
    unknow_1: str
    unknow_2: str
    unknow_3: str
    unknow_4: str
    label: str


def _create_yolo_txt(data: list[GSSDData], dst: Path):
    img_w, img_h = W, H
    sb = Queue()
    for d in data:
        x1, y1 = float(d.x), float(d.y)
        x2, y2 = x1 + float(d.w), y1 + float(d.h)
        x, y = (x1 + x2) / 2 / img_w, (y1 + y2) / 2 / img_h
        w, h = abs(x2 - x1) / img_w, abs(y2 - y1) / img_h
        label_idx = TYPE_DICT[d.label]
        sb.put(f"{label_idx} {x:.4f} {y:.4f} {w:.4f} {h:.4f}")
    
    with dst.open("w") as fp:
        fp.write("\n".join(list(sb.queue)))

def do_create(src: Path, dst: Path):
    if not(src.is_file()) or not(dst.is_dir()):
        raise ValueError("输入文件不是单个文件，或者输出目录不是文件夹")
    
    cache = Queue()
    cur_frame_id = "no frame"
    with src.open("r") as fp:
        for line in fp:
            sp = map(lambda t: t.strip(), line.split(","))
            cur_data = GSSDData._make(sp)

            if cur_data.frame_id != cur_frame_id:
                if not cache.empty():
                    _create_yolo_txt(list(cache.queue), dst / f"{int(cur_frame_id):08}.txt")
                cache.queue.clear()
                cur_frame_id = cur_data.frame_id
            
            cache.put(cur_data)
        
        if not cache.empty():
            _create_yolo_txt(list(cache.queue), dst / f"{int(cur_frame_id):08}.txt")
